import os
os.environ['MINERU_TOOLS_CONFIG_JSON'] = '/home/ly2/magic-pdf.json'
os.environ['CUDA_VISIBLE_DEVICES'] = '2'
from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader
from magic_pdf.data.dataset import PymuDocDataset
from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
from magic_pdf.config.enums import SupportedPdfParseMethod

# args
pdf_file_name = "/data/ly2/rag_ly/pdf_files/《高效使用DeepSeek》.pdf"  # replace with the real pdf path
name_without_suff = '.'.join(os.path.basename(pdf_file_name).split(".")[:-1])
# prepare env
local_image_dir, local_md_dir = "output/images", "output"
image_dir = str(os.path.basename(local_image_dir))
os.makedirs(local_image_dir, exist_ok=True)
image_writer = FileBasedDataWriter(local_image_dir)
md_writer = FileBasedDataWriter(local_md_dir)

# read bytes
reader1 = FileBasedDataReader("")
pdf_bytes = reader1.read(pdf_file_name)  # read the pdf content
ds = PymuDocDataset(pdf_bytes)

if ds.classify() == SupportedPdfParseMethod.OCR:
    infer_result = ds.apply(doc_analyze, ocr=True)
    pipe_result = infer_result.pipe_ocr_mode(image_writer)

else:
    infer_result = ds.apply(doc_analyze, ocr=False)
    pipe_result = infer_result.pipe_txt_mode(image_writer)
# pipe_result.dump_content_list(md_writer, f"{name_without_suff}_content_list.json", image_dir)
infer_result.dump_model(md_writer, f"{name_without_suff}_model.json")
pipe_result.dump_middle_json(md_writer, f"{name_without_suff}_middle.json")
pipe_result.dump_md(md_writer, f"{name_without_suff}.md", image_dir)
